#include "treegen.h"
#include <vector>
#include <stdio.h>
#include <stdlib.h>
#include "shadow2d.h"
#include <QFile>
#include <queue>
#include <assert.h>
#include "constants.h"

#define PARAMNUM 19
#define BUDNUM 44

using namespace std;

//comment for roger


TreeGen::TreeGen(){
    treeRoot = new Bud();
    //_tropism = float3(0,TROPISM,0);
    _wind = float3(0,0,0);
    numBuds = 0;
    //srand(468367346);
     LAMBDA= .58;
    //alpha modifies resource at base
     ALPHA= 2;
    //determines rate of diameter growth (smaller the fatter)
     DIAMETER_RATE= BRANCH_DIAMETER_RATE;
    // >0 for growth influenced up, <0 for growth skewed down
     TROPISM= .4;
    //default angle of auxilary branch
     AUX_ANGLE= 0.75;
    //weight of default aux direction in branch growth
     AUX_WEIGHT= 3;
    //weight of default term direction in branch growth
     TERM_WEIGHT= 3;
    //scale adujst for term and aux directions
     // was 4
     TLENGTH= 2;
     ALENGTH= 2;
     SHED_THRES= .5;
     in_interact = false;
     this->m_addMultMeta = true;
     m_apicalControl = true;
     m_useBH = true;
    branchIdGen = 0;
     m_apicalDepth  = 1;
     W_MAX = 1.0;
     W_MIN = .006;
     KAPPA = 0.5;
}

TreeGen::~TreeGen(){
    freeTree(treeRoot);
}

void TreeGen::freeTree(Bud *root){
    if( root->term)
        freeTree(root->term);
    if( root->aux)
        freeTree(root->aux);
    delete root;

}

//BH MODEL:

//This method gathers resources and collects at the base
void TreeGen::basipetalBH(Bud *b){
    float tq =0;

    //if no aux bud add the resource form auxilary
    if(!b->aux) tq += b->aQ;
    else{
        basipetalBH(b->aux);
        tq += b->aux->iQ;
    }
    if(!b->term) tq += b->tQ;
    else{
        basipetalBH(b->term);
        tq += b->term->iQ;
    }
    b->iQ = tq;
}

//This method should be called to collect resources from the buds
//and redistribute back from base
void TreeGen::passBH(bool shed){
    basipetalBH(treeRoot);
    this->treeRoot->resource = ALPHA * this->treeRoot->iQ;
    pass2BH(this->treeRoot);
    if(shed)
        this->shedBranches(this->treeRoot);
}

//This method redistributes resources back to buds
void TreeGen::pass2BH(Bud *b){
    float qm = b->term ? b->term->iQ : b->tQ;
    float ql = b->aux ? b->aux->iQ : b->aQ;
    if(!b->aux && b->aux_fate == ABORT)
        ql =0;
    if(!b->term && b->term_fate == ABORT)
        qm = 0;

    float qt =  (LAMBDA * qm + (1.0-LAMBDA)* ql);

    float vm = b->resource * LAMBDA * qm / qt;
    float vl =b->resource * (1.0- LAMBDA)* ql / qt;

    if(qt ==0){
        vm = 0;
        vl = 0;
    }

    if(b->aux){
        //This is not part of the algorithm but rather a quick fix to stop it from going downward
        //I think the downward tren is just an error in the algorithm.
        if(b->aux->direction.y <-.8 && !in_interact) {  vl*=.5;}
        b->aux->resource = vl;
        pass2BH(b->aux);
    }
    else{
        b->ares = vl;

    }
    //cout << "ares:" << b->ares << endl;

    if(b->term){
        //This is not part of the algorithm but rather a quick fix to stop it from going downward
        //I think the downward tren is just an error in the algorithm.
        if(b->term->direction.y < -0.8 && !in_interact){ vm*=.5; }
        b->term->resource = vm;
          pass2BH(b->term);
    }
    else{
        b->tres = vm;
    }
}


//PRIORITY MODEL:

void TreeGen::passPriority(bool shed){
    Branch *rootbranch = new Branch();
    rootbranch->startBud = this->treeRoot;
    this->gatherP(treeRoot,rootbranch);
    rootbranch->resource = (ALPHA + 3.0) * rootbranch->Q;
    this->distributeP(rootbranch, 0);
    if(shed)
        this->shedBranchesWithPriority(rootbranch);
    this->clearBranches(rootbranch);

}

void TreeGen::gatherP(Bud *b, Branch *cur){
    bool add_aux = true;
    if(b->term){
        gatherP(b->term, cur);
        cur->numinternodes++;

    }
    else{
       Branch *branch = new Branch();
       branch->just_bud = true;
       branch->Q = b->tQ;
       branch->bud = b;
       branch->bud_is_aux = false;
       branch->numbuds = 1;
       branch->startBud = b;
       cur->tBranch = branch;
       cur->Q += branch->Q;

       //cur->numbuds++;
       cur->numbuds += branch->numbuds;
       add_aux = false;
    }

    if(b->aux){
        Branch *branch = new Branch();
        branch->startBud = b->aux;
        gatherP(b->aux, branch);

        //cur->numbuds++;
        cur->numinternodes += branch->numinternodes;
        cur->numbuds += branch->numbuds;
        cur->Q += branch->Q;
        cur->children.push_back(branch);

    }
    else if(add_aux && b->aux_fate != ABORT){
        Branch *branch = new Branch();
        branch->just_bud = true;
        branch->Q = b->aQ;
        branch->bud = b;
        branch->bud_is_aux = true;
        branch->numbuds = 1;
        branch->startBud = b;

        cur->children.push_back(branch);
        cur->Q += branch->Q;
        cur->numbuds += branch->numbuds;
    }

}

void TreeGen::distributeP(Branch *cur, int depth){
    priority_queue<Branch*, vector<Branch*>, greater<Branch*> > queue;
    //cur->ratio = cur->Q / (float) cur->numbuds;
    cur->tBranch->ratio = cur->tBranch->Q / 1.0;
    bool apical = depth < m_apicalDepth;
    for(int i=0; i< cur->children.size(); i++){
        cur->children[i]->ratio = cur->children[i]->Q / (float) cur->children[i]->numbuds;
    }
    for(int i=0; i< cur->children.size(); i++){
        queue.push(cur->children[i]);
    }

    float ind = 0;
    float sum = 0;

    if(apical){
        cur->tBranch->w = W_MAX;
        cur->tBranch->resource = W_MAX * cur->tBranch->Q;
        sum += cur->tBranch->resource;
        ind++;
    }
    else{
        queue.push(cur->tBranch);
    }
    float w = W_MIN ;
    float N = (float) cur->children.size() + 1.0;
    while(!queue.empty()){
        if(ind > KAPPA * N )
            w = W_MIN;
        else
            w = W_MAX + (W_MIN - W_MAX) / (KAPPA * N) * ind;
        Branch *b = (Branch*) queue.top();
        queue.pop();
        b->resource = w * b->Q;
        sum += b->resource;
        ind+=1.0;
    }

    if(sum == 0) sum = 1;
    assert(cur->tBranch);
    cur->tBranch->bud->tres = cur->tBranch->resource * cur->resource / sum;

    for(int i = 0; i< cur->children.size(); i++){
        cur->children[i]->resource *= cur->resource;
        cur->children[i]->resource /= sum;
        if(cur->children[i]->just_bud){
            if(cur->children[i]->children.size() > 0) cout<<"children's children size should be 0";
            if(cur->children[i]->bud_is_aux){
                cur->children[i]->bud->ares = cur->children[i]->resource;
            }
            else{
                cur->children[i]->bud->tres = cur->children[i]->resource;
            }
        }
        else{
            distributeP(cur->children[i], depth + 1);
        }
    }
}

void TreeGen::shedBranchesWithPriority(Branch *b){
    if((b->resource / b->numinternodes * 60.0) < SHED_THRES && !b->just_bud){
        if(b->startBud->aux)
            freeTree(b->startBud->aux);
        if(b->startBud->term)
            freeTree(b->startBud->term);
        b->startBud->aux = NULL;
        b->startBud->term = NULL;
        b->startBud->term_fate = ABORT;
        b->startBud->aux_fate = ABORT;
    }
    else {
        foreach(Branch* child, b->children){
            if(!((Branch*) child)->just_bud)
                shedBranchesWithPriority((Branch*) child);
        }
    }
}

void TreeGen::clearBranches(Branch *b){
    foreach(Branch* child, b->children){
        clearBranches((Branch*) child);
    }
    delete b;
}



void TreeGen::createBudFromBud(Bud *b, Bud *out){
    out->budsAbove = 0;
    out->start = b->end;
    out->parent = b;
}
//SAVED AT THIS POINT
void TreeGen::produceShootsNew(Bud *b, int depth){
    float ladj = m_addMultMeta ? 0 : 1.5;
    if(! b->aux){
        //create one metmer: old way
        if( b->ares > 1.0 && b->aux_fate != ABORT){
            b->aux = new Bud();
            createBudFromBud(b,b->aux);
            float3 dir = b->auxDir;
            dir.normalize();
            b->aux->direction = dir;

            float l = (ALENGTH + ladj)* b->ares / ((float) (int) b->ares);
            b->aux->length = l;
            b->aux->end = b->aux->length*(b->aux->direction) + b->aux->start;
            b->aux->diameter = 8;
            initializeBud(b->aux);
            if(b->ares > 4.0) b->ares = 4.0;
            //now create more:
            if(m_addMultMeta){
                b->ares -= 1.0;
                Bud *cur = b->aux;
                while(b->ares >1.0){
                    cur->term = new Bud();
                    createBudFromBud(cur, cur->term);
                    float tWeight = depth < m_apicalDepth && !m_useBH ? 3*TERM_WEIGHT : TERM_WEIGHT;
                    float3 dir = tWeight*cur->direction + b->aV + float3(0,TROPISM,0) + (_wind*b->end.y);
                    //dir = cur->direction;
                    dir.normalize();
                    cur->term->direction = dir;
                    //maybe should be same length as first one
                    //cur->term->length = TLENGTH* b->ares / ((float) (int) b->ares);
                    cur->term->length = l;
                    cur->term->end = cur->term->length*(cur->term->direction) + cur->term->start;
                    cur->term->diameter = 8;
                    initializeBud(cur->term);
                    cur = cur->term;
                    b->ares -= 1.0;
                }
                cur->term_fate = FLOWER;
                //b->aux->bud_fate = DORMANT;
            }

        }
    }
    else
        produceShootsNew(b->aux, depth +1);


    if(!b->term){

        if(b->tres > 1.0 && b->term_fate != ABORT){
            b->term = new Bud();
            createBudFromBud(b,b->term);
            float3 dir = b->direction;
            dir.normalize();
            b->term->direction = dir;
            float l = (TLENGTH + ladj) * b->tres / ((float) (int) b->tres);
            b->term->length = l;
            b->term->end = b->term->length*(b->term->direction) + b->term->start ;
            b->term->theta = atan2(b->term->end.y - b->term->start.y,b->term->end.x - b->term->start.x);
            b->term->diameter = 8;
            initializeBud(b->term);
            //clamp max number of metamers
            if(b->tres > 4.0) b->tres = 4.0;

            if(m_addMultMeta){
                b->tres -= 1.0;
                Bud *cur = b->term;
                while(b->tres > 1.0){
                    cur->term = new Bud();
                    createBudFromBud(cur, cur->term);
                    float tWeight = depth < m_apicalDepth && !m_useBH ? 3*TERM_WEIGHT : TERM_WEIGHT;
                    float3 dir = tWeight*cur->direction + b->tV + float3(0,TROPISM,0) + (_wind*b->end.y);
                    //dir = cur->direction;
                    dir.normalize();
                    cur->term->direction = dir;
                    //maybe should be same length as first one

                    //cur->term->length = TLENGTH* b->tres / ((float) (int) b->tres);
                    cur->term->length = l;
                    cur->term->end = cur->term->length*(cur->term->direction) + cur->term->start;
                    cur->term->diameter = 8;
                    initializeBud(cur->term);
                    cur = cur->term;
                    b->tres -= 1.0;
                }
                cur->term_fate = FLOWER;
            }
        }
    }
    else
        produceShootsNew(b->term,depth); //see header for why depth is not incremented
}


void TreeGen::produceShootsOld(Bud *b){
    //bool oneMetamer = !this->m_tMeta;
    //may want to play around with setting one Metamer to false and seeing what happens.
    float ladj = m_addMultMeta ? 0 : 1.5;
    if(! b->aux){
        //create one metmer: old way
        if( b->ares > 1.0 && b->aux_fate != ABORT){
            b->aux = new Bud();
            createBudFromBud(b,b->aux);
            float3 dir = AUX_WEIGHT*b->auxDir + b->aV + float3(0,TROPISM,0) + (_wind*b->end.y);
            dir.normalize();
            b->aux->direction = dir;

            float l = (ALENGTH + ladj)* b->ares / ((float) (int) b->ares);
            b->aux->length = l;
            b->aux->end = b->aux->length*(b->aux->direction) + b->aux->start;
            b->aux->diameter = 8;
            initializeBud(b->aux);

            //now create more:
             if(m_addMultMeta){
                b->ares -= 1.0;
                Bud *cur = b->aux;
                while(b->ares >1.0){
                    cur->term = new Bud();
                    createBudFromBud(cur, cur->term);
                    float3 dir = TERM_WEIGHT*cur->direction + b->aV + float3(0,TROPISM,0) + (_wind*b->end.y);
                    dir = cur->direction;
                    dir.normalize();
                    cur->term->direction = dir;
                    //maybe should be same length as first one
                    //cur->term->length = TLENGTH* b->ares / ((float) (int) b->ares);
                    cur->term->length = l;
                    cur->term->end = cur->term->length*(cur->term->direction) + cur->term->start;
                    cur->term->diameter = 8;
                    initializeBud(cur->term);
                    cur = cur->term;
                    b->ares -= 1.0;
                }
                cur->term_fate = FLOWER;
                //b->aux->bud_fate = DORMANT;
            }

        }
    }
    else
        produceShootsOld(b->aux);


    if(!b->term){
        if(b->tres > 1.0 && b->term_fate != ABORT){
            b->term = new Bud();
            createBudFromBud(b,b->term);
            float3 dir = TERM_WEIGHT*b->direction + b->tV + float3(0,TROPISM,0) + (_wind*b->end.y);
            dir.normalize();
            b->term->direction = dir;
            float l = (TLENGTH + ladj) * b->tres / ((float) (int) b->tres);
            b->term->length = l;
            b->term->end = b->term->length*(b->term->direction) + b->term->start ;
            b->term->theta = atan2(b->term->end.y - b->term->start.y,b->term->end.x - b->term->start.x);
            b->term->diameter = 8;
            initializeBud(b->term);

            if(m_addMultMeta){
                b->tres -= 1.0;
                Bud *cur = b->term;
                while(b->tres > 1.0){
                    cur->term = new Bud();
                    createBudFromBud(cur, cur->term);
                    float3 dir = TERM_WEIGHT*cur->direction + b->tV + float3(0,TROPISM,0) + (_wind*b->end.y);
                    dir = cur->direction;
                    dir.normalize();
                    cur->term->direction = dir;
                    //maybe should be same length as first one

                    //cur->term->length = TLENGTH* b->tres / ((float) (int) b->tres);
                    cur->term->length = l;
                    cur->term->end = cur->term->length*(cur->term->direction) + cur->term->start;
                    cur->term->diameter = 8;
                    initializeBud(cur->term);
                    cur = cur->term;
                    b->tres -= 1.0;
                }
                cur->term_fate = FLOWER;
            }
        }
    }
    else
        produceShootsOld(b->term);
}

void TreeGen::shedBranches(Bud *b){
    int num = 0;
    if(b->aux){
        shedBranches(b->aux);
        num += b->aux->budsAbove;
        num++;
    }    

    if(b->term){
        shedBranches(b->term);
        num+= b->term->budsAbove;
        num++;
    }


    b->budsAbove = num;
    if(num != 0){
        if((b->resource / (double) num) < SHED_THRES){
            if(b->aux)
                freeTree(b->aux);
            if(b->term)
                freeTree(b->term);
            b->aux = NULL;
            b->term = NULL;
            b->term_fate = ABORT;
            b->aux_fate = ABORT;
        }

    }

}


void TreeGen::TreeToEvanFormat(){
   //  j center.x center.y center.z radius 0 0 0 radius 0 0 0 radius parentIndex
   //  (where parentIndex is 1-based not 0-based, so 0 is no parent)
    m_out = new QFile("/home/bcheung/Desktop/savedTree.txt");
    m_out->open(QIODevice::WriteOnly | QIODevice::Text);
    QTextStream out(m_out);
    out<<"j "<< 0 <<" "<< 0<<" " << 0;
    out<<" "<<treeRoot->diameter/2.0 <<" 0 0 0 "<<treeRoot->diameter/2.0<<" 0 0 0 "<<treeRoot->diameter/2.0;
    out<<" 0"<<"\n";

    m_nextIndx =2;
    this->recurSave(this->treeRoot, 1, out);
    m_out->close();
    delete m_out;

}

void TreeGen::recurSave(Bud *bud, int parentindx,  QTextStream &out){
    out<<"j "<< bud->end.x - treeRoot->start.x<<" "<< bud->end.y - treeRoot->start.y<<" "
            << bud->end.z - treeRoot->start.z;
    out<<" "<<bud->diameter/2.0 <<" 0 0 0 "<<bud->diameter/2.0<<" 0 0 0 "<<bud->diameter/2.0;
    out<<" "<<parentindx<<"\n";
    int myindx = m_nextIndx;
    if(bud->aux){
        m_nextIndx++;
        recurSave(bud->aux, myindx, out);
    }
    if(bud->term){
        m_nextIndx++;
        recurSave(bud->term, myindx, out);
    }

}













vector <float *> TreeGen::treeToArray()
{
        vector <float *> v;
        int l = 0;
        if (treeRoot == NULL)
            cout << "TreeRoot NULL " << endl;
        float * Array = new float[PARAMNUM];
        Array[0] = m_nextIndx;
        Array[1] = in_interact ;
        Array[2] = m_addMultMeta ;
        Array[3] = m_apicalControl ;
        Array[4] = m_useBH ;
        Array[5] = LAMBDA ;
        Array[6] = ALPHA ;
        Array[7] = DIAMETER_RATE ;
        Array[8] = TROPISM ;
        Array[9] = AUX_ANGLE ;
        Array[10] = AUX_WEIGHT ;
        Array[11] = TERM_WEIGHT ;
        Array[12] = TLENGTH ;
        Array[13] = ALENGTH ;
        Array[14] = SHED_THRES ;
        Array[15] = W_MAX ;
        Array[16] = W_MIN ;
        Array[17] = KAPPA ;
        Array[18] = numBuds ;

        v.push_back(Array);
        toArrayRecur(v,treeRoot,0, l);
        cout << "arrLength:" << l << "----" << v.size()-1 << endl;
        return v;
}



void TreeGen::toArrayRecur(vector <float *> & v, Bud * bud, int index, int & arrLength)
{
    float * Array = new float[BUDNUM];
    Array[ 0] = 0; //index;
    Array[ 1] = bud->term_fate;
    Array[ 2] = bud->aux_fate;
    Array[ 3] = bud->budsAbove;
    Array[ 4] = bud->exposure;
    Array[ 5] = bud->darken;
    Array[ 6] = bud->rotleaf;
    Array[ 7] = bud->season;
    Array[ 8] = bud->direction.x;
    Array[ 9] = bud->direction.y;
    Array[ 10] = bud->direction.z;
    Array[ 11] = bud->diameter;
    Array[ 12] = bud->length;
    Array[ 13] = bud->phi;
    Array[ 14] = bud->theta;
    Array[ 15] = bud->start.x;
    Array[ 16] = bud->start.y;
    Array[ 17] = bud->start.z;
    Array[ 18] = (bud->leaf) ? 1.0 : 0.0;
    Array[ 19] = bud->i;
    Array[ 20] = bud->j;
    Array[ 21] = bud->k;
    Array[ 22] = bud->tQ;
    Array[ 23] = bud->aQ;
    Array[ 24] = bud->iQ;
    Array[ 25] = bud->resource;
    Array[ 26] = bud->ares;
    Array[ 27] = bud->tres;
    Array[ 28] = bud->auxAngle;
    Array[ 29] = bud->auxDir.x;
    Array[ 30] = bud->auxDir.y;
    Array[ 31] = bud->auxDir.z;
    Array[ 32] = bud->tV.x;
    Array[ 33] = bud->tV.y;
    Array[ 34] = bud->tV.z;
    Array[ 35] = bud->aV.x;
    Array[ 36] = bud->aV.y;
    Array[ 37] = bud->aV.z;
    Array[ 38] = bud->spaceColIndex;
    Array[ 39] = bud->end.x;
    Array[ 40] = bud->end.y;
    Array[ 41] = bud->end.z;
    Array[ 42] = (bud->term) ? bud->term->spaceColIndex : -1;
    Array[ 43] = (bud->aux) ? bud->aux->spaceColIndex : -1;
    arrLength++;
    v.push_back(Array);
    if (bud->aux)
    {
        //cout << "auxaux" << endl;
        toArrayRecur(v,bud->aux,0,arrLength);
    }
    if (bud->term)
    {
        //cout << "termterm" << endl;
        toArrayRecur(v,bud->term,0,arrLength);
    }

}

vector <float*> TreeGen::readFromFile( QString name)
{
        vector <float*> v;
        QFile * min = new QFile(name);
        if (!min->open(QIODevice::ReadOnly | QIODevice::Text))
        {
            cout << "Error: Cannot read from file: " << name.toStdString()<< ". Please check the name again." << endl;
            return v;
        }
        QTextStream in(min);

        QString line = in.readLine();
        QStringList firstLine = line.split(" ");
        float * array = new float [PARAMNUM];

        for (int i = 0; i < PARAMNUM; i ++)
        {
            QString curr = firstLine.at(i);
            array[i] = curr.toFloat();
        }
        v.push_back(array);

        line = in.readLine();
        while (!line.isNull())
        {
            firstLine = line.split(" ");
            float * array = new float[BUDNUM];
            for (int i = 0; i < BUDNUM; i++)
            {
                QString curr = firstLine.at(i);
                array[i] = curr.toFloat();
            }
            v.push_back(array);
            line = in.readLine();
        }
        cout << "end: " << v.size() << endl;
        return v;
        min->close();
}

bool TreeGen::writeToFile(vector <float*> v, QString name)
{
    cout << "real value: " << v.size() << endl;
    int l = v.size();
    QFile * mout = new QFile(name);
    if (!mout->open(QIODevice::WriteOnly | QIODevice::Text))
    {
        cout << "Error: Cannot open file: " << name.toStdString() << ". Please check the name again." << endl;
        return false;
    }

    QTextStream out(mout);

    float * first = v.at(0);

    for (int i = 0; i < PARAMNUM; i++)
    {
        out << first[i] << " ";
    }
    out << endl;

    for (int i = 1; i < l; i++)
    {
        float * Array = v.at(i);

        for (int j = 0; j < BUDNUM; j++)
            out << Array[j] << " " ;
        out << endl;
    }

    mout->close();
    return true;
}


bool TreeGen::saveTreeToFile( QString filename)
{
        vector <float * > arr = treeToArray();
        return writeToFile(arr,filename);
}

